import { z } from 'zod'

import type { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js'

interface RequiredEnv {
	AI: Ai
}

// Zod schema for AI Search response validation
const AiSearchResponseSchema = z.object({
	object: z.string(),
	search_query: z.string(),
	data: z.array(
		z.object({
			file_id: z.string(),
			filename: z.string(),
			score: z.number(),
			attributes: z
				.object({
					modified_date: z.number().optional(),
					folder: z.string().optional(),
				})
				.catchall(z.any()),
			content: z.array(
				z.object({
					id: z.string(),
					type: z.string(),
					text: z.string(),
				})
			),
		})
	),
	has_more: z.boolean(),
	next_page: z.string().nullable(),
})

/**
 * Registers the docs search tool with the MCP server using AI Search
 * @param server The MCP server instance
 */
export function registerDocsTools(server: McpServer, env: RequiredEnv) {
	server.tool(
		'search_cloudflare_documentation',
		`Search the Cloudflare documentation.

		This tool should be used to answer any question about Cloudflare products or features, including:
		- Workers, Pages, R2, Images, Stream, D1, Durable Objects, KV, Workflows, Hyperdrive, Queues
		- AI Search, Workers AI, Vectorize, AI Gateway, Browser Rendering
		- Zero Trust, Access, Tunnel, Gateway, Browser Isolation, WARP, DDOS, Magic Transit, Magic WAN
		- CDN, Cache, DNS, Zaraz, Argo, Rulesets, Terraform, Account and Billing

		Results are returned as semantically similar chunks to the query.
		`,
		{
			query: z.string(),
		},
		{
			title: 'Search Cloudflare docs',
			annotations: {
				readOnlyHint: true,
			},
		},
		async ({ query }) => {
			const results = await queryAiSearch(env.AI, query)
			const resultsAsXml = results
				.map((result) => {
					return `<result>
<url>${result.url}</url>
<title>${result.title}</title>
<text>
${result.text}
</text>
</result>`
				})
				.join('\n')
			return {
				content: [{ type: 'text', text: resultsAsXml }],
			}
		}
	)

	// Note: this is a tool instead of a prompt because
	// prompt support is much less common than tools.
	server.tool(
		'migrate_pages_to_workers_guide',
		`ALWAYS read this guide before migrating Pages projects to Workers.`,
		{},
		{
			title: 'Get Pages migration guide',
			annotations: {
				readOnlyHint: true,
			},
		},
		async () => {
			const res = await fetch(
				'https://developers.cloudflare.com/workers/prompts/pages-to-workers.txt',
				{
					cf: { cacheEverything: true, cacheTtl: 3600 },
				}
			)

			if (!res.ok) {
				return {
					content: [{ type: 'text', text: 'Error: Failed to fetch guide. Please try again.' }],
				}
			}

			return {
				content: [
					{
						type: 'text',
						text: await res.text(),
					},
				],
			}
		}
	)
}

async function queryAiSearch(ai: Ai, query: string) {
	const rawResponse = await doWithRetries(() =>
		ai.autorag('docs-mcp-rag').search({
			query,
		})
	)

	// Parse and validate the response using Zod
	const response = AiSearchResponseSchema.parse(rawResponse)

	return response.data.map((item) => ({
		similarity: item.score,
		id: item.file_id,
		url: sourceToUrl(item.filename),
		title: extractTitle(item.filename),
		text: item.content.map((c) => c.text).join('\n'),
	}))
}

function sourceToUrl(filename: string): string {
	// Convert filename to URL format
	// Example: "workers/configuration/index.md" -> "https://developers.cloudflare.com/workers/configuration/"
	return (
		'https://developers.cloudflare.com/' +
		filename.replace(/index\.mdx?$/, '').replace(/\.mdx?$/, '')
	)
}

function extractTitle(filename: string): string {
	// Extract a reasonable title from the filename
	// Example: "workers/configuration/index.md" -> "Configuration"
	const parts = filename.replace(/\.mdx?$/, '').split('/')
	const lastPart = parts[parts.length - 1]

	if (lastPart === 'index') {
		// Use the parent directory name if filename is index
		return parts[parts.length - 2] || 'Documentation'
	}

	// Convert kebab-case or snake_case to title case
	return lastPart.replace(/[-_]/g, ' ').replace(/\b\w/g, (l) => l.toUpperCase())
}

/**
 * Retries an action with exponential backoff, only for retryable errors
 * @template T
 * @param {() => Promise<T>} action
 */
async function doWithRetries<T>(action: () => Promise<T>) {
	const NUM_RETRIES = 5
	const INIT_RETRY_MS = 100

	for (let i = 0; i <= NUM_RETRIES; i++) {
		try {
			return await action()
		} catch (e) {
			// Check if error is retryable (system errors, not user errors)
			const isRetryable = isRetryableError(e)

			console.error(`AI Search attempt ${i + 1} failed:`, e)

			if (!isRetryable || i === NUM_RETRIES) {
				throw e
			}

			// Exponential backoff with jitter
			const delay = Math.random() * INIT_RETRY_MS * Math.pow(2, i)
			await scheduler.wait(delay)
		}
	}
	// Should never reach here – last loop iteration should throw
	throw new Error('An unknown error occurred')
}

/**
 * Determines if an error is retryable based on error type and status
 */
function isRetryableError(error: unknown): boolean {
	// Handle HTTP errors from fetch-like responses
	if (error && typeof error === 'object' && 'status' in error) {
		const status = (error as { status: number }).status
		// Retry server errors (5xx) and rate limits (429), not client errors (4xx)
		return status >= 500 || status === 429
	}

	// Handle network errors, timeouts, etc.
	if (error instanceof Error) {
		const errorMessage = error.message.toLowerCase()
		return (
			errorMessage.includes('timeout') ||
			errorMessage.includes('network') ||
			errorMessage.includes('connection') ||
			errorMessage.includes('fetch')
		)
	}

	// Default to retryable for unknown errors (conservative approach)
	return true
}
